# 4、使用光流估计方法，在前述测试视频上计算特征点，进一步进行特征点光流估计。
import cv2 as cv

videoFileName = r"D:\AI\Picture\vtest.avi"
# 角点检测参数
feature_params = dict(maxCorners=100,
                      qualityLevel=0.3,
                      minDistance=7,
                      blockSize=7)
# lucas kanade 光流法参数
lk_params = dict(winSize=(15, 15),
                 maxLevel=2,
                 criteria=(cv.TERM_CRITERIA_EPS | cv.TERM_CRITERIA_COUNT, 10, 0.03))
# 加载视频
cap = cv.VideoCapture(videoFileName)

# 计算第一帧特征点

# ret表示是否成功获取帧, prev是捕获到的第一帧图像
ret, prev = cap.read()
# 第一帧图像灰度化
prevGray = cv.cvtColor(prev, cv.COLOR_BGR2GRAY)
# 使用 Shi-Tomasi 方法获取图像中 N 个最好的角点（也可以通过改变参数来使用 Harris 角点检测算法）。
# 通常情况下，输入的应该是灰度图像
p0 = cv.goodFeaturesToTrack(prevGray, mask=None, **feature_params)

while True:
    # ret表示是否成功获取帧, frame是捕获到的图像
    ret, frame = cap.read()
    # 没有读到当前帧，结束
    if not ret:
        break
    # 灰度图像
    gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
    """
    使用具有金字塔的迭代Lucas-Kanade方法计算稀疏特征集的光流。
    返回值：   
    p1 输出一个二维点的向量，这个向量可以是用来作为光流算法的输入特征点，也是光流算法在当前帧找到特征点的新位置（浮点数）
    st 标志，在当前帧当中发现的特征点标志st==1，否则为0
    err 向量中的每个特征对应的错误率
    """
    p1, st, err = cv.calcOpticalFlowPyrLK(prevGray, gray, p0, None, **lk_params)

    # 选取好的跟终点,goodPoints和goodPrevPoints都是点的列表集
    goodPoints = p1[st == 1]
    goodPrevPoints = p0[st == 1]

    # 在结果图像中迭加画出特征点和计算出来的光流向量
    res = frame.copy()
    drawColor = (0, 0, 255)
    # zip() 函数用于将可迭代的对象作为参数，将对象中对应的元素打包成一个个元组，然后返回由这些元组组成的列表。
    # enumerate() 函数用于将一个可遍历的数据对象(如列表、元组或字符串)组合为一个索引序列，同时列出数据和数据下标。
    # ravel() 是用来将矩阵向量化的操作
    for i, (cur, prev) in enumerate(zip(goodPoints, goodPrevPoints)):
        x0, y0 = cur.ravel()
        x1, y1 = prev.ravel()
        cv.line(res, (x0, y0), (x1, y1), drawColor)
        cv.circle(res, (x0, y0), 3, drawColor)

    # 更新上一帧
    prevGray = gray.copy()
    # reshape()是数组array中的方法，作用是将数据重新组织,这里改为一个3维数组
    p0 = goodPoints.reshape(-1, 1, 2)

    # 显示计算结果图像
    cv.imshow("result", res)

    # 每一帧间隔30ms
    key = cv.waitKey(30)
    if key == 27:
        break

cap.release()
cv.destroyAllWindows()
